# ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv
# ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv
import warnings
warnings.filterwarnings('ignore')

import os,pdb
import pandas as pd
import numpy as np
from util_set_zh_matplot import plt
import seaborn as sns
import networkx as nx
import shap
import joblib

from sklearn.preprocessing import LabelEncoder, StandardScaler, PolynomialFeatures
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix
from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import (
    RandomForestClassifier, 
)

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

# 创建输出目录
out_dir = 'ques4_prepare_data'
os.makedirs(out_dir, exist_ok=True)

'''
# ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv
# ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv
分别在 染色体的非整倍体(包含大概七个类别) 绘制柱状图，显示每个类别的数量

两个统计量相加，再显示一个整体的类别柱状图

'''
def main():
    # prepare_data()
    stage1()

def prepare_data():
    # 加载数据
    female_data = pd.read_csv('ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv')
    male_data = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

    # 去除所有的 unname 列
    # 1. 去除所有的未命名列
    def remove_unnamed_columns(df):
        # 删除所有以Unnamed开头的列
        unnamed_cols = [col for col in df.columns if 'Unnamed' in col]
        df = df.drop(columns=unnamed_cols, errors='ignore')
        return df
    
    female_data = remove_unnamed_columns(female_data)
    male_data = remove_unnamed_columns(male_data)

    # 2. 列名标准化处理
    # 统一"唯一比对的读段数"列名
    male_data = male_data.rename(columns={'唯一比对的读段数  ': '唯一比对的读段数'})
    
    # 3. 处理不一致字段
    # 女胎特有列
    female_specific_cols = ['孕周天数']
    female_data = female_data.drop(columns=female_specific_cols, errors='ignore')
    
    # 男胎特有列
    male_specific_cols = ['实际孕周', '达标孕周', '达标', 'Y染色体的Z值', 'Y染色体浓度']
    male_data = male_data.drop(columns=male_specific_cols, errors='ignore')
    
    # 4. 添加性别列
    female_data['性别'] = '女'
    male_data['性别'] = '男'

    # 5. 合并数据集
    combined_data = pd.concat([female_data, male_data], axis=0, ignore_index=True)
    
    # 6. 保存合并后的数据
    combined_data.to_csv(os.path.join(out_dir , 'ques4_concat.csv'), 
                         index=False, encoding='utf-8-sig')
    
    print("数据已成功合并并保存为 ques4_concat.csv")
    print(f"总记录数: {len(combined_data)} (女胎: {len(female_data)}, 男胎: {len(male_data)})")
    print("合并后的列名:", combined_data.columns.tolist())
    
    # return combined_data

    # 验证两个数据的列名相互一致
    # 2. 验证两个数据的列名相互一致
    def validate_columns(df1, df2):
        cols1 = set(df1.columns)
        cols2 = set(df2.columns)
        
        # 检查列名是否完全相同
        if cols1 == cols2:
            print("列名验证通过：两个数据集具有完全相同的列名")
            return True
        else:
            # 找出差异
            only_in_female = cols1 - cols2
            only_in_male = cols2 - cols1
            
            if only_in_female:
                print(f"仅在女胎数据中存在的列：{only_in_female}")
            if only_in_male:
                print(f"仅在男胎数据中存在的列：{only_in_male}")
            
            return False

    # 保存回原来的位置
    is_valid = validate_columns(female_data, male_data)
    
    # 3. 保存回原来的位置
    if is_valid:
        print("列名一致")
    else:
        print("列名不一致，请先处理列名差异后再保存")

    female_counts = female_data['染色体的非整倍体'].value_counts()
    male_counts = male_data['染色体的非整倍体'].value_counts()

    # 合并两个统计量
    total_counts = female_counts.add(male_counts, fill_value=0)

    print(female_counts)
    print(male_counts)
    print(total_counts)

    plt.figure(figsize=(15, 10))

    # 女胎数据柱状图
    plt.subplot(3, 1, 1)
    female_counts.plot(kind='bar', color='pink')
    plt.title('女胎染色体非整倍体类别统计')
    plt.ylabel('数量')

    # 男胎数据柱状图
    plt.subplot(3, 1, 2)
    male_counts.plot(kind='bar', color='lightblue')
    plt.title('男胎染色体非整倍体类别统计')
    plt.ylabel('数量')

    # 总体数据柱状图
    plt.subplot(3, 1, 3)
    total_counts.plot(kind='bar', color='purple')
    plt.title('总体染色体非整倍体类别统计')
    plt.ylabel('数量')

    plt.tight_layout()
    # plt.show()
    plt.savefig( os.path.join( out_dir , 'ques4_染色体非整倍体类别统计.png' ) )
    plt.close()

def stage1():
    file_path = 'ques4_prepare_data\ques4_concat.csv'
    combined_data = pd.read_csv(file_path)
    # female_data male_date 从 combined_data 分离出来
    # 2. 从合并数据中分离出女胎和男胎数据
    female_data = combined_data[combined_data['性别'] == '女'].copy()
    male_data = combined_data[combined_data['性别'] == '男'].copy()


    # 2. 设置阈值过滤稀有类别（样本数<10）
    value_counts = combined_data['染色体的非整倍体'].value_counts()
    valid_categories = value_counts[value_counts >= 10].index
    filtered_data = combined_data[combined_data['染色体的非整倍体'].isin(valid_categories)]

    print(f"原始类别分布:\n{value_counts}\n")
    print(f"过滤后保留的类别:\n{valid_categories}\n")

    # 3. 确保男女标签一致（取交集）
    female_labels = set(female_data['染色体的非整倍体'].unique())
    male_labels = set(male_data['染色体的非整倍体'].unique())
    common_labels = list(female_labels & male_labels)
    common_labels = [label for label in common_labels if label in valid_categories]  # 结合阈值过滤

    final_data = filtered_data[filtered_data['染色体的非整倍体'].isin(common_labels)]
    print(f"最终使用的共同类别:\n{common_labels}\n")

    # 4. 定义特征列（根据您提供的列表）
    features = [
        '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', '唯一比对的读段数',
        'GC含量', '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
        'X染色体的Z值', 'X染色体浓度', '13号染色体的GC含量', 
        '18号染色体的GC含量', '21号染色体的GC含量'
    ]

    # 检查特征是否存在
    missing_features = [f for f in features if f not in final_data.columns]
    if missing_features:
        raise ValueError(f"以下特征在数据中不存在: {missing_features}")

    # 5. 数据准备
    X = final_data[features]
    y = final_data['染色体的非整倍体']
    gender = final_data['性别']


    # 打印平衡前分布
    print("平衡前类别分布:")
    print(y.value_counts())

    # 6. 数据平衡处理 - 将每个类别样本量限制为700
    balanced_data = pd.DataFrame()
    for g in ['女', '男']:
        for label in common_labels:
            # 获取当前性别和类别的数据
            mask = (gender == g) & (y == label)
            subset = final_data[mask].copy()
            
            # 如果样本量超过700，随机抽取700个
            if len(subset) > 700:
                subset = subset.sample(n=700, random_state=42)
            # 如果样本量不足700，复制样本直到达到700
            elif len(subset) < 700:
                repeat_times = (700 // len(subset)) + 1
                subset = pd.concat([subset]*repeat_times, ignore_index=True).sample(n=700, random_state=42)
            
            balanced_data = pd.concat([balanced_data, subset], ignore_index=True)



    print("\n平衡后类别分布:")
    print(balanced_data['染色体的非整倍体'].value_counts())

    # 6. 数据标准化（在训练集上拟合，应用到测试集）
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(balanced_data[features])
    balanced_data_scaled = pd.DataFrame(X_scaled, columns=features)
    balanced_data_scaled['染色体的非整倍体'] = balanced_data['染色体的非整倍体'].values
    balanced_data_scaled['性别'] = balanced_data['性别'].values

    # 7. 交叉验证方案
    le = LabelEncoder()
    y_encoded = le.fit_transform(balanced_data_scaled['染色体的非整倍体'])

    # 案例1：用男性训练，女性测试
    X_train = balanced_data_scaled[balanced_data_scaled['性别'] == '男'][features]
    y_train = y_encoded[balanced_data_scaled['性别'] == '男']
    X_test = balanced_data_scaled[balanced_data_scaled['性别'] == '女'][features]
    y_test = y_encoded[balanced_data_scaled['性别'] == '女']

    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)

    print("\n男性训练集预测女性测试集结果:")
    print(classification_report(y_test, y_pred, target_names=le.classes_))

    # 案例2：用女性训练，男性测试
    X_train = balanced_data_scaled[balanced_data_scaled['性别'] == '女'][features]
    y_train = y_encoded[balanced_data_scaled['性别'] == '女']
    X_test = balanced_data_scaled[balanced_data_scaled['性别'] == '男'][features]
    y_test = y_encoded[balanced_data_scaled['性别'] == '男']


    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)

    print("\n女性训练集预测男性测试集结果:")
    print(classification_report(y_test, y_pred, target_names=le.classes_))

    # 8. 特征重要性分析
    print("\n特征重要性排序:")
    importance = pd.DataFrame({
        'feature': features,
        'importance': clf.feature_importances_
    }).sort_values('importance', ascending=False)
    print(importance)
if __name__=='__main__':
    main()